import os
import PIL.Image as PImage
from PIL import ImageFont, ImageDraw
import cv2
import json
import numpy as np
import torch
import torch.utils.data as data
import scipy.misc
import torchvision.transforms as transforms
try:
    from Tkinter import *
    from ttk import *
    from tkFileDialog import *
    from tkMessageBox import *
except ImportError:
    from tkinter import *
    from tkinter.ttk import *
    from tkinter.filedialog import *
    from tkinter.messagebox import *
import random
import sys
sys.path.append('/home/chengk/tool-repo/idcardgenerator')
from idcardgenerator import gui
import skimage.measure
from ..utils.transforms import generate_target, transform_pixel


# 获取人名素材
import pandas as pd
df_id = pd.read_excel('/home/chengk/params-photo-serving/chk/身份证.xlsx', )
names = df_id['transact_name'].tolist()
characters = np.unique(list(''.join(names))).tolist()
id_number = df_id['id_number'].tolist()

def keepLargestConnection(mask):
    labels, num_labels = skimage.measure.label(mask>0, return_num=True)
    label_most_cnt, label_most = 0, -1
    for label in range(1, num_labels+1):
        label_cnt = (labels == label).sum()
        if label_cnt > label_most_cnt:
            label_most = label
            label_most_cnt = label_cnt
    if label_most_cnt > 0:
        mask[(labels!=label_most)] = 0
    return mask

import glob
def portrait_generate():
    portrait_file = np.random.choice(glob.glob('/home/chengk/tool-repo/idcardgenerator/portraits/*.png'), 1)[0]
    # 500x670
    img = PImage.open(portrait_file)
    w, h = img.size
    ratio = min(w/500, h/670)
    w_, h_ = int(500*ratio), int(670*ratio)
    img = img.crop(((w-w_)//2, (h-h_)//2, (w-w_)//2+w_, (h-h_)//2+h_))
    img = img.resize((500, 670))
    return img

def id_number_generate():
    return np.random.choice(id_number, 1).tolist()[0]

def gender_generate():
    if np.random.rand() >= 0.5:
        return '男'
    else:
        return '女'

def birth_generate(birth):
    year = birth[6: 10]
    month = birth[10:12]
    day = birth[12:14]
    return year, month, day

def name_generate():
    # 3字、2字、4字的概率分别为0.8, 0.1, 0.1
    dice = np.random.rand()
    if dice <= 0.8:
        return ''.join(np.random.choice(characters, 3, replace=True))
    elif dice <= 0.9:
        return ''.join(np.random.choice(characters, 2, replace=True))
    else:
        return ''.join(np.random.choice(characters, 4, replace=True))

def address_generate():
    char_size = 8+int(np.random.rand()*15)
    return ''.join(np.random.choice(characters, char_size))

def mask_idcard(img):
    mask = (img[..., 0]==img[..., 1])&(img[..., 0]==img[..., 2])&(img[..., 0]==0)
    mask = 1-keepLargestConnection(mask)
    return mask

def rand(min_val=0, max_val=1):
    while True:
        val = np.random.rand()
        if val >= min_val and val <= max_val:
            return val

upper_left, upper_right, bottom_left, bottom_right = (289, 490), (2170, 494), (284, 1677), (2174, 1679)
def standard_front_side():
    y_split = 1830
    base_dir = '/home/chengk/tool-repo/idcardgenerator/idcardgenerator/usedres/'
    # name, sex, nation, year, mon, day, addr, idn, org, life, bgvar
    
    nation = '汉'
    addr = address_generate()
    org, life, bgvar = '广州公安局', 'gzpolice', '2020-1-1'
    sex = gender_generate()
    name = name_generate()
    idn = id_number_generate()
    year, mon, day = birth_generate(idn)
    
    im = PImage.open(os.path.join(base_dir, 'empty.png'))
    avatar = portrait_generate()

    name_font = ImageFont.truetype(os.path.join(base_dir, 'hei.ttf'), 72)
    other_font = ImageFont.truetype(os.path.join(base_dir, 'hei.ttf'), 60)
    bdate_font = ImageFont.truetype(os.path.join(base_dir, 'fzhei.ttf'), 60)
    id_font = ImageFont.truetype(os.path.join(base_dir, 'ocrb10bt.ttf'), 72)

    draw = ImageDraw.Draw(im)
    draw.text((630, 690), name, fill=(0, 0, 0), font=name_font)
    draw.text((630, 840), sex, fill=(0, 0, 0), font=other_font)
    draw.text((1030, 840), nation, fill=(0, 0, 0), font=other_font)
    draw.text((630, 980), year, fill=(0, 0, 0), font=bdate_font)
    draw.text((950, 980), mon, fill=(0, 0, 0), font=bdate_font)
    draw.text((1150, 980), day, fill=(0, 0, 0), font=bdate_font)
    start = 0
    loc = 1120
    while start + 11 < len(addr):
        draw.text((630, loc), addr[start:start + 11], fill=(0, 0, 0), font=other_font)
        start += 11
        loc += 100
    draw.text((630, loc), addr[start:], fill=(0, 0, 0), font=other_font)
    draw.text((950, 1475), idn, fill=(0, 0, 0), font=id_font)
    draw.text((1050, 2750), org, fill=(0, 0, 0), font=other_font)
    draw.text((1050, 2895), life, fill=(0, 0, 0), font=other_font)

    avatar = cv2.cvtColor(np.asarray(avatar), cv2.COLOR_RGBA2BGRA)
    im = cv2.cvtColor(np.asarray(im), cv2.COLOR_RGBA2BGRA)
    im = gui.changeBackground(avatar, im, (500, 670), (690, 1500))
    im = PImage.fromarray(cv2.cvtColor(im, cv2.COLOR_BGRA2RGBA))
    im = im.crop((0, 0, im.size[0], y_split))
    im = im.crop((0, 0, im.size[0], im.size[1]+(490-(im.size[1]-1679))))
    mask = mask_idcard(np.array(im))
    im = np.array(im)
    im[mask>0, -1] = 255
    im[mask==0, -1] = 0
    im = PImage.fromarray(im)
    
    return im, mask

def paste_a2b(img_rgba_a, img_rgba_b):
    '''
        img_rgba_a: 前景图片 4通道 uint8类型
        img_rgba_b: 背景图片 4通道 uint8类型
        返回：
            img_rgba_a_over_b: 叠放图片 4通道 uint8类型
    '''
    img_rgba_a = img_rgba_a.astype(np.float)
    img_rgba_b = img_rgba_b.astype(np.float)
    # 透明度压缩到0~1
    img_rgba_a[..., -1] /= 255
    img_rgba_b[..., -1] /= 255

    img_rgba_a_over_b = img_rgba_a.copy()
    img_rgba_a_over_b[..., -1] = img_rgba_a[..., -1] + img_rgba_b[..., -1] * (1-img_rgba_a[..., -1]) 
    img_rgba_a_over_b[..., :3] = (img_rgba_a[..., :3]*img_rgba_a[..., 3:] + img_rgba_b[..., :3] * img_rgba_b[..., 3:] * (1-img_rgba_a[..., 3:])) / img_rgba_a_over_b[..., -1:]
    img_rgba_a_over_b[..., -1] *= 255
    return img_rgba_a_over_b.clip(0, 255).astype(np.uint8)

# def crop_square(img):
#     '''
#     img: PIL.Image
#     '''
#     w, h = img.size
#     bounds = np.array([0, 0, w, h])
#     if h > w:
#         padding = np.array([(w-h)//2, 0, h-w+(w-h)//2, 0])
#     else:
#         padding = np.array([0, (h-w)//2, 0, w-h+(h-w)//2])
#     bounds += padding
    
#     img = img.crop(bounds.tolist())
#     return img

def extract_anno(anno):
    anno = anno[0]
    return np.array([[anno['content'][0]['x'], anno['content'][0]['y']],
                     [anno['content'][1]['x'], anno['content'][1]['y']],
                     [anno['content'][3]['x'], anno['content'][3]['y']],
                     [anno['content'][2]['x'], anno['content'][2]['y']]], dtype=np.float32)

def crop_square(img, return_pad=False):
    '''
    img: PIL.Image
    '''
    w, h = img.size
    assert img.size[0] == w
    assert img.size[1] == h
    bounds = np.array([0, 0, w, h])
    if h > w:
        padding = np.array([(w-h)//2, 0, h-w+(w-h)//2, 0])
    else:
        padding = np.array([0, (h-w)//2, 0, w-h+(h-w)//2])
    bounds += padding
    
    img = img.crop(bounds.tolist())
    assert img.size[0] == img.size[1]
    if return_pad:
        return img, np.abs(padding)
    else:
        return img

class IDCard(data.Dataset):
    """idcard
    """
    def __init__(self, cfg, is_train=True, transform=None, length=2000):
        self.is_train = is_train
        self.transform = transform
        self.mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
        self.std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
        self.output_size = cfg.MODEL.HEATMAP_SIZE
        self.input_size = cfg.MODEL.IMAGE_SIZE
        self.sigma = cfg.MODEL.SIGMA
        self.label_type = cfg.MODEL.TARGET_TYPE

        self.bg_crop = transforms.Compose([crop_square,
            transforms.RandomResizedCrop(self.input_size)])
        self.resize_square = transforms.Resize(self.input_size)
        self.color_jitter = transforms.ColorJitter(brightness=0.3, contrast=0.2, saturation=0.2, hue=.2)
        self.bg_folder = '/home/chengk/tool-repo/idcardgenerator/clothing-pattern-dataset/FingerCamera/'
        self.bg_set = []
        # 固定背景集合
        for folder in os.listdir(self.bg_folder):
            self.bg_set += glob.glob(os.path.join(self.bg_folder, folder, '*.png'))
        # print(self.bg_set)
        self.img_set = [(f, '') for f in glob.glob('/home/chengk/tool-repo/idcardgenerator/idcard_set/*.png')]
        self.img_set += [('./real_idcard/'+f.rsplit('/', 1)[-1].replace('json', 'png'), f) for f in glob.glob('./real_anno/*.json')]

        print('现有身份证%d张'%len(self.img_set))
        if self.is_train:
            self.length = len(self.img_set)
        else:
            self.length = 200

    def __len__(self):
        return self.length

    def random_combine(self, bg, img, loc):
        '''
            bg: PIL.Image
            img: PIL.Image

            return: PIL.Image
        '''
        # 图片变换
        bg = self.bg_crop(bg)
        mask = PImage.fromarray(np.array(img)[..., -1])
        img = self.color_jitter(img)
        # 转换具体坐标为比例
        loc = loc.astype(np.float) / np.array(img.size)
        
        w_idcard, h_idcard = img.size
        resize_ratio = self.input_size[0]/max(w_idcard, h_idcard)
        # Here comes the random part
        resize_ratio *= rand(0.5)
        img = img.resize((int(w_idcard*resize_ratio), int(h_idcard*resize_ratio))) # 透明度丢失
        mask = mask.resize((int(w_idcard*resize_ratio), int(h_idcard*resize_ratio)))
        # 重新贴合透明度层
        img = np.array(img)
        img[..., -1] = np.array(mask)
        
        w_idcard, h_idcard = img.shape[1::-1]
        
        x_start = np.random.randint(self.input_size[0]-w_idcard)
        y_start = np.random.randint(self.input_size[0]-h_idcard)
        # 计算当前顶点坐标
        loc = loc * np.array([w_idcard, h_idcard])
        loc += np.array([x_start, y_start])
        
        palette = np.zeros((self.input_size[0], self.input_size[0], 4), dtype=np.uint8)
        palette[y_start: y_start+h_idcard, x_start: x_start+w_idcard, :] = img
        
        return PImage.fromarray(paste_a2b(palette, np.array(bg))), loc


    def __getitem__(self, idx):
        # img, _ = standard_front_side() # 生成固定身份证照片
        fn, anno = self.img_set[idx]
        img = PImage.open(fn)
        if anno == '':
            loc = np.array([upper_left, upper_right, bottom_left, bottom_right]) # pre-defined
            bg = PImage.open(np.random.choice(self.bg_set, 1)[0])
        else:
            with open(anno, 'r') as f:
                anno_json = json.load(f)
            loc = extract_anno(anno_json)
        if anno == '':
            img_combined, loc = self.random_combine(bg, img, loc)
        else:
            img, padding = crop_square(img, True)
            loc += padding[:2][None, ...]
            loc *= self.input_size[0] / img.size[0]
            img = self.resize_square(img)
            img = self.color_jitter(img)
            img_combined = img
        # print('img_combined.shape:', img_combined.size)

        nparts = loc.shape[0]
        # print('顶点:', loc)
        # 构造heatmap
        target = np.zeros((nparts, self.output_size[0], self.output_size[1]))
        tpts = loc.copy()/4
        center = np.mean(tpts, -1)
        # print('center:', center)
        for i in range(nparts):
            # print('scaled loc:', tpts[i])
            target[i] = generate_target(target[i], tpts[i], self.sigma, label_type=self.label_type)

        img_combined = np.array(img_combined)[..., 2::-1].astype(np.float32)
        img_combined = (img_combined/255.0 - self.mean) / self.std
        img_combined = img_combined.transpose([2, 0, 1])

        target = torch.Tensor(target)
        meta = {'center': torch.Tensor(center), 'scale': 1, 'index': idx, 'pts': loc}

        return img_combined, target, meta